Skip to content

fix: preserve streamed AFC function call chunks#2453

Open
MukundaKatta wants to merge 1 commit into
googleapis:mainfrom
MukundaKatta:codex/stream-afc-merge-function-call-chunks
Open

fix: preserve streamed AFC function call chunks#2453
MukundaKatta wants to merge 1 commit into
googleapis:mainfrom
MukundaKatta:codex/stream-afc-merge-function-call-chunks

Conversation

@MukundaKatta
Copy link
Copy Markdown
Contributor

Fixes #1938.\n\n## Summary\n- merge streamed AFC function-call chunks before appending them to the next request\n- preserve thought signatures for parallel tool calls split across stream chunks\n- add sync and async regression coverage for multi-chunk function calls\n\n## Tests\n- uv run --with pytest --with pytest-asyncio --with mcp pytest google/genai/tests/afc/test_generate_content_stream_afc.py -q\n\nNote: I also tried the broader chat send_message suite, but it requires replay/API credentials in this environment and failed before exercising this patch.

@edsponsler
Copy link
Copy Markdown

edsponsler commented May 20, 2026

Thanks for putting this fix together!

I pulled down this branch (refs/pull/2453/head) locally to test it against the function calling issues we've been seeing with Gemini 3.1.

The Good News: This PR successfully resolves the issue for streaming workloads! When testing with chat.send_message_stream() (which uses generate_content_stream/async_generator under the hood), the multi-turn function calling executes perfectly without dropping the thought_signature. This fully resolves the crash reported in #1938.

The Remaining Issue (Ref: #2406): The PR appears to be incomplete. While it fixes the chunk-merging logic for streams, it does not fix the synchronous generate_content method in models.py.

Because standard blocking calls like chat.send_message() route through the unpatched generate_content logic, they are still crashing on the subsequent remote tool call with the exact same error: 400 INVALID_ARGUMENT: Function call is missing a thought_signature in functionCall parts.

Here is a minimal reproduction script using google-genai==2.3.0 installed directly from this PR branch. It consistently crashes on the synchronous execution.

import os
from google import genai
from google.genai import types

def get_files_info(directory: str = ".") -> str:
    print(f" -> Executing tool: get_files_info(directory='{directory}')")
    return "file1.txt, main.py"

def get_file_content(file_path: str) -> str:
    print(f" -> Executing tool: get_file_content(file_path='{file_path}')")
    return "print('Hello world')"

# Fails on PR branch #2453 with: "Function call is missing a thought_signature..."
def main():
    # Make sure to set GEMINI_API_KEY in your environment before running
    client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))

    chat = client.chats.create(
        model="gemini-3.1-flash-lite",
        config=types.GenerateContentConfig(
            tools=[get_files_info, get_file_content],
            automatic_function_calling=types.AutomaticFunctionCallingConfig(
                disable=False, maximum_remote_calls=5
            ),
            # Prevents parallel execution to force sequential multi-turn AFC
            system_instruction="Do not call tools in parallel. Always list files first.",
            temperature=0.0,
        ),
    )

    prompt = "Please look at the files in the current directory, and then tell me what is inside the python file."
    
    try:
        # Standard synchronous send_message (routes to unpatched `generate_content`)
        response = chat.send_message(prompt)
        print("\n--- SUCCESS ---")
        print(response.text)
    except Exception as e:
        print("\n--- CRASH DETECTED ---")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()

If you modify this script to use response = chat.send_message_stream(prompt), it succeeds flawlessly.

Could we get a similar patch applied to the generate_content AFC loop (around line 5000 in models.py) so that the synchronous clients are fixed before this merges?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

AsyncChat.send_message_stream does not pass thought_signature with function calls

2 participants